from pickletools import optimize
from myconv import Conv
import torch
import torch.nn as nn
import time

device = "cuda"
# 生成测试数据
input_data = torch.randn([1, 3, 33, 6]).to(device)
gt = torch.randn([1, 1, 33, 6]).to(device)

# 自制卷积与 torch 卷积，参数都初始化为 0.5
conv = Conv(in_channels=3, out_channels=1, kernel_size=3, stride=1, padding=1).to(device)
conv_torch = torch.nn.Conv2d(in_channels=3, out_channels=1, kernel_size=3, stride=1, padding=1).to(device)
for param in conv_torch.parameters():
    nn.init.constant_(param, 0.5)

num = 30000
time_start_torch = time.time()
# Torch 方法
for i in range(num):
    # 前向传播
    res_conv_torch = conv_torch(input_data)

    # 建立优化器与损失
    opt4torch = torch.optim.Adam(conv_torch.parameters())
    loss_fn = torch.nn.MSELoss()

    # 反向传播
    opt4torch.zero_grad()

    loss4torch = loss_fn(gt, res_conv_torch)

    loss4torch.backward()

    opt4torch.step()
time_end_torch = time.time()

time_start_rolfma = time.time()
# 自写方法
for i in range(num):
    # 前向传播
    res_conv = conv(input_data)

    # 建立优化器与损失
    opt4rolfma = torch.optim.Adam(conv.parameters())
    loss_fn = torch.nn.MSELoss()

    # 反向传播
    opt4rolfma.zero_grad()

    loss4rolfma = loss_fn(gt, res_conv)

    loss4rolfma.backward()

    opt4rolfma.step()
time_end_rolfma = time.time()

print(res_conv - res_conv_torch)
print("torch time:", time_end_torch - time_start_torch)
print("rolfma time:", time_end_rolfma - time_start_rolfma)
print("rate:", (time_end_rolfma - time_start_rolfma) / (time_end_torch - time_start_torch))
